package cn.sciento.starter.keyencrypt.mvc;

import cn.sciento.core.convertor.ApplicationContextHelper;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Type;
import java.nio.charset.StandardCharsets;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import javax.servlet.http.HttpServletResponse;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.collections4.MapUtils;
import org.apache.commons.lang3.StringUtils;
import cn.sciento.starter.keyencrypt.core.Encrypt;
import cn.sciento.starter.keyencrypt.core.EncryptContext;
import cn.sciento.starter.keyencrypt.core.IEncryptionService;
import cn.sciento.starter.keyencrypt.util.EncryptUtils;
import org.springframework.beans.BeanUtils;
import org.springframework.core.MethodParameter;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.lang.Nullable;
import org.springframework.web.accept.ContentNegotiationManager;
import org.springframework.web.bind.support.WebDataBinderFactory;
import org.springframework.web.context.request.NativeWebRequest;
import org.springframework.web.context.request.ServletWebRequest;
import org.springframework.web.method.support.ModelAndViewContainer;
import org.springframework.web.servlet.mvc.method.annotation.RequestResponseBodyMethodProcessor;

public class EncryptRequestResponseBodyMethodProcessor extends RequestResponseBodyMethodProcessor {
  protected ObjectMapper mapper;
  
  protected IEncryptionService encryptionService;
  
  public EncryptRequestResponseBodyMethodProcessor(List<HttpMessageConverter<?>> converters) {
    super(converters);
  }
  
  public EncryptRequestResponseBodyMethodProcessor(List<HttpMessageConverter<?>> converters, ContentNegotiationManager manager) {
    super(converters, manager);
  }
  
  public EncryptRequestResponseBodyMethodProcessor(List<HttpMessageConverter<?>> converters, List<Object> requestResponseBodyAdvice) {
    super(converters, requestResponseBodyAdvice);
  }
  
  public EncryptRequestResponseBodyMethodProcessor(List<HttpMessageConverter<?>> converters, ContentNegotiationManager manager, List<Object> requestResponseBodyAdvice) {
    super(converters, manager, requestResponseBodyAdvice);
  }
  
  public Object resolveArgument(MethodParameter parameter, @Nullable ModelAndViewContainer mavContainer, NativeWebRequest webRequest, @Nullable WebDataBinderFactory binderFactory) throws Exception {
    if (!EncryptContext.isEncrypt())
      return super.resolveArgument(parameter, mavContainer, webRequest, binderFactory); 
    Encrypt encrypt = parameter.getParameterAnnotation(Encrypt.class);
    if (encrypt == null)
      return super.resolveArgument(parameter, mavContainer, webRequest, binderFactory); 
    if (Collection.class.isAssignableFrom(parameter.getParameterType()))
      return resolveCollectionArgument(parameter, mavContainer, webRequest, binderFactory, encrypt); 
    if (parameter.getParameterType().isArray())
      return resolveArrayArgument(parameter, mavContainer, webRequest, binderFactory, encrypt); 
    if (Map.class.isAssignableFrom(parameter.getParameterType()))
      return resolveMapArgument(parameter, mavContainer, webRequest, binderFactory, encrypt); 
    return super.resolveArgument(parameter, mavContainer, webRequest, binderFactory);
  }
  
  private Object resolveArrayArgument(MethodParameter parameter, @Nullable ModelAndViewContainer mavContainer, NativeWebRequest webRequest, @Nullable WebDataBinderFactory binderFactory, Encrypt encrypt) throws Exception {
    return encryptMultiArgument(parameter, mavContainer, webRequest, binderFactory, encrypt, parameter.getParameterType().getComponentType());
  }
  
  private Object encryptMultiArgument(MethodParameter parameter, @Nullable ModelAndViewContainer mavContainer, NativeWebRequest webRequest, @Nullable WebDataBinderFactory binderFactory, Encrypt encrypt, Class clz) throws Exception {
    if (isSimple(clz)) {
      HttpServletRequest httpServletRequest = (HttpServletRequest)webRequest.getNativeRequest(HttpServletRequest.class);
      HttpServletResponse httpServletResponse = (HttpServletResponse)webRequest.getNativeRequest(HttpServletResponse.class);
      String body = getBodyFromRequest(httpServletRequest);
      if (StringUtils.isBlank(body)) {
        WrapperWebRequest wrapperWebRequest1 = new WrapperWebRequest((HttpServletRequest)new ModifyBodyHttpServletRequestWrapper(httpServletRequest, ""), httpServletResponse);
        return super.resolveArgument(parameter, mavContainer, (NativeWebRequest)wrapperWebRequest1, binderFactory);
      } 
      List<Object> encryptList = (List<Object>)getMapper().readValue(body, List.class);
      List<Object> decryptList = Collections.emptyList();
      if (!CollectionUtils.isEmpty(encryptList))
        decryptList = encryptList.stream().map(t -> (t instanceof String && !EncryptUtils.ignoreValue(encrypt, (String)t)) ? this.encryptionService.decrypt((String)t, encrypt.value(), encrypt.ignoreUserConflict()) : t).collect(Collectors.toList());
      WrapperWebRequest wrapperWebRequest = new WrapperWebRequest((HttpServletRequest)new ModifyBodyHttpServletRequestWrapper(httpServletRequest, getMapper().writeValueAsString(decryptList)), httpServletResponse);
      return super.resolveArgument(parameter, mavContainer, (NativeWebRequest)wrapperWebRequest, binderFactory);
    } 
    return super.resolveArgument(parameter, mavContainer, webRequest, binderFactory);
  }
  
  private Object resolveMapArgument(MethodParameter parameter, @Nullable ModelAndViewContainer mavContainer, NativeWebRequest webRequest, @Nullable WebDataBinderFactory binderFactory, Encrypt encrypt) throws Exception {
    Type type = parameter.getGenericParameterType();
    if (type instanceof ParameterizedType) {
      Type[] params = ((ParameterizedType)type).getActualTypeArguments();
      if (params.length != 2)
        return super.resolveArgument(parameter, mavContainer, webRequest, binderFactory); 
      boolean decryptKey = isSimple((Class)params[0]);
      boolean decryptValue = isSimple((Class)params[1]);
      if (decryptKey || decryptValue) {
        HttpServletRequest httpServletRequest = (HttpServletRequest)webRequest.getNativeRequest(HttpServletRequest.class);
        HttpServletResponse httpServletResponse = (HttpServletResponse)webRequest.getNativeRequest(HttpServletResponse.class);
        String body = getBodyFromRequest(httpServletRequest);
        if (StringUtils.isBlank(body))
          return super.resolveArgument(parameter, mavContainer, webRequest, binderFactory); 
        Map<Object, Object> encryptMap = (Map<Object, Object>)getMapper().readValue(body, Map.class);
        if (MapUtils.isEmpty(encryptMap))
          return super.resolveArgument(parameter, mavContainer, webRequest, binderFactory); 
        Map<Object, Object> decryptMap = new HashMap<>(encryptMap.size());
        encryptMap.forEach((k, v) -> {
              if (decryptKey && k instanceof String && !EncryptUtils.ignoreValue(encrypt, (String)k) && this.encryptionService.isCipher((String)k))
                k = this.encryptionService.decrypt((String)k, encrypt.value(), encrypt.ignoreUserConflict()); 
              if (decryptValue && v instanceof String && !EncryptUtils.ignoreValue(encrypt, (String)v) && this.encryptionService.isCipher((String)v))
                v = this.encryptionService.decrypt((String)v, encrypt.value(), encrypt.ignoreUserConflict()); 
              decryptMap.put(k, v);
            });
        WrapperWebRequest wrapperWebRequest = new WrapperWebRequest((HttpServletRequest)new ModifyBodyHttpServletRequestWrapper(httpServletRequest, getMapper().writeValueAsString(decryptMap)), httpServletResponse);
        return super.resolveArgument(parameter, mavContainer, (NativeWebRequest)wrapperWebRequest, binderFactory);
      } 
    } 
    return super.resolveArgument(parameter, mavContainer, webRequest, binderFactory);
  }
  
  private Object resolveCollectionArgument(MethodParameter parameter, @Nullable ModelAndViewContainer mavContainer, NativeWebRequest webRequest, @Nullable WebDataBinderFactory binderFactory, Encrypt encrypt) throws Exception {
    Type type = parameter.getGenericParameterType();
    if (type instanceof ParameterizedType) {
      Type[] params = ((ParameterizedType)type).getActualTypeArguments();
      if (params.length > 1)
        return super.resolveArgument(parameter, mavContainer, webRequest, binderFactory); 
      return encryptMultiArgument(parameter, mavContainer, webRequest, binderFactory, encrypt, (Class)params[0]);
    } 
    return super.resolveArgument(parameter, mavContainer, webRequest, binderFactory);
  }
  
  public void setEncryptionService(IEncryptionService encryptionService) {
    this.encryptionService = encryptionService;
  }
  
  public ObjectMapper getMapper() {
    if (this.mapper != null)
      return this.mapper; 
    synchronized (this) {
      if (this.mapper != null)
        return this.mapper; 
      this.mapper = ApplicationContextHelper.getContext().getBean(ObjectMapper.class);
    } 
    return this.mapper;
  }
  
  protected String getBodyFromRequest(HttpServletRequest request) {
    ServletInputStream servletInputStream =null;
    if (request == null)
      return null; 
    StringBuilder sb = new StringBuilder();
    InputStream inputStream = null;
    BufferedReader reader = null;
    try {
      servletInputStream = request.getInputStream();
      reader = new BufferedReader(new InputStreamReader((InputStream)servletInputStream, StandardCharsets.UTF_8));
      String line = "";
      while ((line = reader.readLine()) != null)
        sb.append(line); 
    } catch (IOException iOException) {
    
    } finally {
      if (servletInputStream != null)
        try {
          servletInputStream.close();
        } catch (IOException iOException) {} 
      if (reader != null)
        try {
          reader.close();
        } catch (IOException iOException) {} 
    } 
    return sb.toString();
  }
  
  protected boolean isSimple(Class clz) {
    return (BeanUtils.isSimpleProperty(clz) || Object.class.equals(clz));
  }
  
  public class WrapperWebRequest extends ServletWebRequest {
    public WrapperWebRequest(HttpServletRequest request) {
      super(request);
    }
    
    public WrapperWebRequest(HttpServletRequest request, HttpServletResponse response) {
      super(request, response);
    }
  }
  
  public class ModifyBodyHttpServletRequestWrapper extends HttpServletRequestWrapper {
    private String bodyJsonStr;
    
    public ModifyBodyHttpServletRequestWrapper(HttpServletRequest request, String bodyJsonStr) {
      super(request);
      this.bodyJsonStr = bodyJsonStr;
    }
    
    public ServletInputStream getInputStream() throws IOException {
      if (StringUtils.isEmpty(this.bodyJsonStr))
        this.bodyJsonStr = ""; 
      final ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(this.bodyJsonStr.getBytes(StandardCharsets.UTF_8));
      ServletInputStream servletInputStream = new ServletInputStream() {
          public boolean isFinished() {
            return false;
          }
          
          public boolean isReady() {
            return false;
          }
          
          public void setReadListener(ReadListener readListener) {}
          
          public int read() throws IOException {
            return byteArrayInputStream.read();
          }
        };
      return servletInputStream;
    }
    
    public BufferedReader getReader() throws IOException {
      return new BufferedReader(new InputStreamReader((InputStream)getInputStream()));
    }
    
    public String getBodyJsonStr() {
      return this.bodyJsonStr;
    }
    
    public void setBodyJsonStr(String bodyJsonStr) {
      this.bodyJsonStr = bodyJsonStr;
    }
  }
}
